"""
FiftyOne delegated operation repository document.

| Copyright 2017-2025, Voxel51, Inc.
| `voxel51.com <https://voxel51.com/>`_
|
"""
import copy
import logging
from datetime import datetime

from fiftyone.operators.executor import (
    ExecutionContext,
    ExecutionResult,
    ExecutionRunState,
    ExecutionProgress,
)
from fiftyone.operators.types import Pipeline, PipelineRunInfo

logger = logging.getLogger(__name__)


class DelegatedOperationDocument(object):
    def __init__(
        self,
        operator: str = None,
        delegation_target: str = None,
        context: ExecutionContext = None,
        is_remote: bool = False,
    ):
        self.operator = operator
        self.label = None
        self.delegation_target = delegation_target
        if isinstance(context, ExecutionContext) or context is None:
            pass
        else:
            raise AttributeError(
                "context must be an instance of ExecutionContext"
            )
        self.context = context
        self.run_state = (
            ExecutionRunState.SCHEDULED
            if is_remote
            else ExecutionRunState.QUEUED
        )  # if running locally use QUEUED otherwise SCHEDULED
        self.run_link = None
        self.queued_at = datetime.utcnow() if not is_remote else None
        self.updated_at = datetime.utcnow()
        self.status = None
        self.dataset_id = None
        self.started_at = None
        self.pinned = False
        self.completed_at = None
        self.failed_at = None
        self.scheduled_at = datetime.utcnow() if is_remote else None
        self.result = None
        self.id = None
        self._doc = None
        self.metadata = None
        self.log_upload_error = None
        self.log_size = None
        self.log_path = None

        # distributed task fields
        self.parent_id = None  # Only on children
        self.pipeline = None  # Only on pipeline parent
        self.pipeline_run_info = None  # Only on pipeline parent

    @property
    def num_distributed_tasks(self):
        """Returns the number of distributed tasks in this operation, if any."""
        # No distributed tasks
        return None

    def from_pymongo(self, doc: dict):
        # required fields
        self.operator = doc.get("operator")
        self.queued_at = doc.get("queued_at")
        self.run_state = doc.get("run_state")

        # optional fields
        self.delegation_target = doc.get("delegation_target", None)
        self.started_at = doc.get("started_at", None)
        self.completed_at = doc.get("completed_at", None)
        self.failed_at = doc.get("failed_at", None)
        self.scheduled_at = doc.get("scheduled_at", None)
        self.pinned = doc.get("pinned", None)
        self.dataset_id = doc.get("dataset_id", None)
        self.run_link = doc.get("run_link", None)
        self.log_upload_error = doc.get("log_upload_error", None)
        self.log_size = doc.get("log_size", None)
        self.log_path = doc.get("log_path", None)
        self.metadata = doc.get("metadata", None)
        self.label = doc.get("label", None)
        self.updated_at = doc.get("updated_at", None)

        # grouped fields
        self.parent_id = doc.get("parent_id", None)

        # internal fields
        self.id = doc.get("_id", doc.get("id"))
        self._doc = doc

        # nested fields
        if (
            "context" in doc
            and doc["context"] is not None
            and "request_params" in doc["context"]
        ):
            self.context = ExecutionContext(
                request_params=doc["context"]["request_params"],
            )

        if "result" in doc and doc["result"] is not None:
            res = ExecutionResult()
            if "result" in doc["result"]:
                res.result = doc["result"]["result"]
            if "error" in doc["result"]:
                res.error = doc["result"]["error"]

            if res.result or res.error:
                self.result = res

        if "status" in doc and doc["status"] is not None:
            self.status = ExecutionProgress()
            if "progress" in doc["status"]:
                self.status.progress = doc["status"]["progress"]
            if "label" in doc["status"]:
                self.status.label = doc["status"]["label"]
            if "updated_at" in doc["status"]:
                self.status.updated_at = doc["status"]["updated_at"]

        self.pipeline = Pipeline.from_json(doc.get("pipeline"))
        self.pipeline_run_info = PipelineRunInfo.from_json(
            doc.get("pipeline_run_info")
        )

        return self

    def to_pymongo(self) -> dict:
        # We make a copy of self.__dict__ so that changes we make below do not
        # affect the actual object. We exclude certain keys that we don't want
        # to serialize directly. "context" is particularly important we do not
        # try to copy because it may contain big, complicated, non-serializable
        # objects that may cause issues with copying.

        ignore_keys = {
            "_doc",
            "id",
            "context",
            "pipeline",
            "pipeline_run_info",
        }
        d = {
            k: copy.deepcopy(v)
            for k, v in self.__dict__.items()
            if k not in ignore_keys
        }
        if self.context:
            d["context"] = {
                "request_params": self.context._get_serialized_request_params()
            }
        if self.pipeline:
            d["pipeline"] = self.pipeline.to_json()
        if self.pipeline_run_info:
            d["pipeline_run_info"] = self.pipeline_run_info.to_json()

        return d
